import matplotlib
matplotlib.use('Agg')
import numpy as np
import cv2
import imutils
import dlib
from imutils import face_utils
from src.interface import DrowsyDetector
from src.utils import merge_const
from src.gui.settings import USE_MODEL_CONF
import matplotlib.pyplot as plt

merge_const(USE_MODEL_CONF)


def opencv2skimage(image):
    new_image = image.copy()
    new_image[:, :, 0] = image[:, :, 2]
    new_image[:, :, 2] = image[:, :, 0]
    return new_image


class FrameParser(object):

    def __init__(self):
        self.detector = dlib.get_frontal_face_detector()
        self.drowsy_detector = DrowsyDetector()
        self.fatigue_count = 0
        self.previous_show = 'normal'
        self.previous_prob_str = '1.0'

    def get_att_opencv(self, att, size=128):
        fig, ax = plt.subplots()
        plt.xticks([])
        plt.yticks([])
        ax.imshow(att, cmap=plt.cm.afmhot_r)
        fig.canvas.draw()
        att_numpy = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep='')
        att_numpy = att_numpy.reshape(fig.canvas.get_width_height()[::-1] + (3,))
        x = att_numpy[..., 0]
        h, w = x.shape
        for i in range(h):
            if x[i, :].min() != 255:
                l_y = i
                break
        for i in range(h - 1, -1, -1):
            if x[i, :].min() != 255:
                r_y = i
                break
        for j in range(w):
            if x[:, j].min() != 255:
                l_x = j
                break
        for j in range(w - 1, -1, -1):
            if x[:, j].min() != 255:
                r_x = j
                break
        att_numpy = att_numpy[l_y:r_y + 1, l_x:r_x + 1, :]
        att_opencv = cv2.cvtColor(att_numpy, cv2.COLOR_RGB2BGR)
        att_opencv = cv2.resize(att_opencv, (128, 128))
        plt.close()
        return att_opencv

    def crop_face_and_print(self, image):
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        rects = self.detector(gray, 0)
        # 找最大人脸
        max_area = 0
        max_candidate = None
        for (i, rect) in enumerate(rects):
            if hasattr(rect, 'rect'):
                rect = rect.rect
            (x, y, w, h) = face_utils.rect_to_bb(rect)
            if w * h > max_area:
                max_area = w * h
                max_candidate = (x, y, w, h)
        if max_candidate is None:
            return image, None
        x, y, w, h = max_candidate
        x -= 10
        y -= 10
        w += 20
        h += 20
        if x < 0:
            x = 0
        if y < 0:
            y = 0
        cropped = image[y:y + h, x:x + w]
        cropped = imutils.resize(cropped, 128, 128)
        image = cv2.rectangle(image, (x, y), (x + w, y + h), (0, 0, 255), 2)
        return image, cropped

    def parse(self, frame):
        h, w = frame.shape[:2]
        frame = cv2.resize(frame, (int(360 / h * w), 360))
        h, w = frame.shape[:2]
        # 检测人脸
        frame, cropped = self.crop_face_and_print(frame)
        if cropped is not None:
            cropped_RGB = opencv2skimage(cropped)
            pred_label, prob, att = self.drowsy_detector.classify(cropped_RGB)
            prob_str = '%.2f' % prob
            if pred_label == 'normal':
                self.fatigue_count = 0
                show = 'normal'
            elif pred_label == 'yawn':
                self.fatigue_count = 0
                show = 'yawn'
            elif pred_label == 'fatigue':
                self.fatigue_count += 1
                # 稳定fatigue
                if self.fatigue_count >= 1:
                    show = 'fatigue'
                else:
                    show = self.previous_show  # 之前的
                    prob_str = self.previous_prob_str
            if show == 'normal':
                frame = cv2.circle(frame, (30, 30), 12, (0, 255, 0), -1)
                cv2.putText(frame, 'NORMAL', (50, 40), cv2.FONT_HERSHEY_COMPLEX, 1.0, (0, 255, 0), 2)
                cv2.putText(frame, 'prob: ' + prob_str, (18, 65), cv2.FONT_ITALIC, 0.6, (0, 255, 0), 2)
            if show == 'yawn':
                frame = cv2.circle(frame, (30, 30), 12, (0, 255, 255), -1)
                cv2.putText(frame, 'YAWN', (50, 40), cv2.FONT_HERSHEY_COMPLEX, 1.0, (0, 255, 255), 2)
                cv2.putText(frame, 'prob: ' + prob_str, (18, 65), cv2.FONT_ITALIC, 0.6, (0, 255, 255), 2)
            if show == 'fatigue':
                frame = cv2.circle(frame, (30, 30), 12, (0, 0, 255), -1)
                cv2.putText(frame, 'FATIGUE', (50, 40), cv2.FONT_HERSHEY_COMPLEX, 1.0, (0, 0, 255), 2)
                cv2.putText(frame, 'prob: ' + prob_str, (18, 65), cv2.FONT_ITALIC, 0.6, (0, 0, 255), 2)
            self.previous_show = show
            self.previous_prob_str = prob_str
            # display att
            if att is not None:
                att_opencv = self.get_att_opencv(att)
                frame[h - 128:h, 0:128, :] = att_opencv
        return frame
